﻿
#include "media_session.h"
#include <initguid.h>
#include <windows.h>
#include <mfapi.h>
#include <mfidl.h>
#include <evr.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <wincodec.h>

// 需要 mf.lib mfplat.lib windowscodecs.lib

#ifdef _MSC_VER
#pragma comment(lib, "mf.lib")
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mfuuid.lib")
#pragma comment(lib, "strmiids.lib")
#pragma comment(lib, "windowscodecs.lib")
#endif


#define SAFE_RELEASE(comobj) if(comobj)comobj->lpVtbl->Release(comobj)

static int init(){
    static int inited = 0;
    if(inited)return 0;
    CoInitializeEx(NULL,0);
    MFStartup(MF_VERSION, MFSTARTUP_LITE);    
    inited = 1;
    return 0;
}

static int media_device_list(const GUID* type, media_device_list_callback cb, void* udata){
    HRESULT hr;
    IMFAttributes* pAttributes = NULL;
    IMFActivate** ppDevices = NULL; 
    UINT32 i, count, len;
    LPWSTR wname = NULL;
    char name[256];
    int ret = -1;
    count = 0;  
    if(init())return -1;
    if(!cb)return -2;
    hr = MFCreateAttributes(&pAttributes, 1); if(FAILED(hr))return -1;    
    hr = pAttributes->lpVtbl->SetGUID(pAttributes, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, type); if(FAILED(hr))goto clean;    
    hr = MFEnumDeviceSources(pAttributes, &ppDevices, &count); if(FAILED(hr))goto clean;    
    for(i = 0; i < count; i++){
        hr = ppDevices[i]->lpVtbl->GetAllocatedString(ppDevices[i], &MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &wname, &len); if(FAILED(hr))goto clean;        
        name[0] = 0;
        WideCharToMultiByte(CP_UTF8, 0, wname, -1, name, 256, NULL, NULL);
        cb(name, udata);
        CoTaskMemFree(wname);
        wname = NULL;
    }    
    ret = 0;   
clean:
    for(i = 0; i < count; i++){
        SAFE_RELEASE(ppDevices[i]);
    }    
    SAFE_RELEASE(pAttributes);
    if(ppDevices)CoTaskMemFree(ppDevices);
    if(wname)CoTaskMemFree(wname);    
    return ret;    
}

int video_input_device_list(media_device_list_callback cb, void* udata){    
    return media_device_list(&MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID, cb, udata);
}

int audio_input_device_list(media_device_list_callback cb, void* udata){
    return media_device_list(&MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_AUDCAP_GUID, cb, udata);
}


struct media_session{
    IMFMediaSession* session;
    IMFTopology* topology;
    IMFVideoDisplayControl* control;
    bool first;
};

typedef struct media_input{
    IMFTopologyNode* node;
    IMFMediaSource* source; 
    IMFPresentationDescriptor* pd;
    IMFStreamDescriptor* sd;   
}media_input;

typedef struct media_output{
    IMFTopologyNode* node;
    IMFActivate* act;
    //IMFMediaSink* sink;    
}media_output;

void media_session_destroy(media_session_t m){
    if(m){
        SAFE_RELEASE(m->session);
        SAFE_RELEASE(m->topology);
        SAFE_RELEASE(m->control);
    }
}
void video_input_destory(video_input_t v){
    media_input* in = (media_input*)v;
    if(in){
        SAFE_RELEASE(in->node);
        SAFE_RELEASE(in->source);
        SAFE_RELEASE(in->pd);
        SAFE_RELEASE(in->sd);
    }    
}
void video_output_destory(video_output_t v){
    media_output* out = (media_output*)v;
    if(out){
        SAFE_RELEASE(out->node);
        SAFE_RELEASE(out->act);
    }    
}

void audio_input_destory(audio_input_t v){
    media_input* in = (media_input*)v;
    if(in){
        SAFE_RELEASE(in->node);
        SAFE_RELEASE(in->source);
        SAFE_RELEASE(in->pd);
        SAFE_RELEASE(in->sd);
    }      
}
void audio_output_destory(audio_output_t v){
    media_output* out = (media_output*)v;
    if(out){
        SAFE_RELEASE(out->node);
        SAFE_RELEASE(out->act);
    }    
}

media_session_t media_session_create(){
    HRESULT hr;
    media_session_t ret = malloc(sizeof(struct media_session));
    if(!ret)return ret;
    if(init())return NULL;
    ret->session = NULL;
    ret->topology = NULL;
    hr = MFCreateMediaSession(NULL, &ret->session); if(FAILED(hr))goto error;
    hr = MFCreateTopology(&ret->topology); if(FAILED(hr))goto error;
    ret->first = true;
    ret->control = NULL;
    //hr = ret->session->lpVtbl->SetTopology(ret->session, MFSESSION_SETTOPOLOGY_IMMEDIATE, ret->topology); if(FAILED(hr))goto error;
    return ret;
error:
    SAFE_RELEASE(ret->session);
    SAFE_RELEASE(ret->topology);
    if(ret)free(ret);
    return NULL;
}


static int media_session_add_input(media_session_t m, media_input* in){
    HRESULT hr;
    if(!m || !in)return -1;
    if(!in->node){
        hr = MFCreateTopologyNode(MF_TOPOLOGY_SOURCESTREAM_NODE, &in->node); if(FAILED(hr))goto error;
        hr = in->node->lpVtbl->SetUnknown(in->node, &MF_TOPONODE_SOURCE, (IUnknown*)in->source); if(FAILED(hr))goto error;
        hr = in->node->lpVtbl->SetUnknown(in->node, &MF_TOPONODE_PRESENTATION_DESCRIPTOR, (IUnknown*)in->pd); if(FAILED(hr))goto error;
        hr = in->node->lpVtbl->SetUnknown(in->node, &MF_TOPONODE_STREAM_DESCRIPTOR, (IUnknown*)in->sd); if(FAILED(hr))goto error;
    }
    hr = m->topology->lpVtbl->AddNode(m->topology, in->node); if(FAILED(hr))goto error;
    return 0;
error:
    return -1;    
}

static int media_session_add_output(media_session_t m, media_output* out){
    HRESULT hr;  
    if(!m || !out)return -1;  
    if(!out->node){
        hr = MFCreateTopologyNode(MF_TOPOLOGY_OUTPUT_NODE, &out->node); if(FAILED(hr))goto error;       
        hr = out->node->lpVtbl->SetObject(out->node, (IUnknown*)out->act);  if(FAILED(hr))goto error;        
    }
    hr = m->topology->lpVtbl->AddNode(m->topology, out->node); if(FAILED(hr))goto error;
    return 0;
error:    
    return -1;     
}

static int media_connect(media_input* in, media_output* out){
    HRESULT hr = in->node->lpVtbl->ConnectOutput(in->node, 0, out->node, 0);
    if(hr)return -1;
    return 0;
}

int media_session_add_video(media_session_t m, video_input_t in, video_output_t out){
    int ret = 0;
    if(in)ret |= media_session_add_input(m, (media_input*)in);
    if(out)ret |= media_session_add_output(m, (media_output*)out);    
    if(in && out)ret |= media_connect((media_input*)in, (media_output*)out);    
    return ret;
}

int media_session_add_audio(media_session_t m, audio_input_t in, audio_output_t out){
    int ret = 0;
    if(in)ret |= media_session_add_input(m, (media_input*)in);
    if(out)ret |= media_session_add_output(m, (media_output*)out);
    if(in && out)ret |= media_connect((media_input*)in, (media_output*)out);   
    return ret;    
}

int media_session_start(media_session_t m){
    HRESULT hr;
    PROPVARIANT p;
    if(m->first){
        hr = m->session->lpVtbl->SetTopology(m->session, 0, m->topology); if(FAILED(hr))goto error; 
        m->first = false; 
    }  
    PropVariantClear(&p); 
    hr = m->session->lpVtbl->Start(m->session, NULL, &p);
    return FAILED(hr)?-1:0;
error:
    return -1;    
}

int media_session_stop(media_session_t m){
    HRESULT hr;
    hr = m->session->lpVtbl->Stop(m->session);
    return FAILED(hr)?-1:0;    
}


static int media_input_init(media_input* in){
    HRESULT hr;
    IMFPresentationDescriptor *pPD = NULL;
    IMFStreamDescriptor *pSD = NULL;
    IMFMediaSource* pSource = in->source;
    BOOL fSelected;
    hr = pSource->lpVtbl->CreatePresentationDescriptor(pSource, &pPD); if (FAILED(hr))goto error;     
    hr = pPD->lpVtbl->GetStreamDescriptorByIndex(pPD, 0, &fSelected, &pSD);  if(FAILED(hr))goto error;
    in->pd = pPD;
    in->sd = pSD;       
    return 0;
error:
    return -1;    
}

static int media_input_create(media_input* in, const char* devname, const GUID* type){
    HRESULT hr;
    IMFAttributes* pAttributes = NULL;
    IMFActivate** ppDevices = NULL; 
    UINT32 i, count, len;
    LPWSTR wname = NULL;
    int ret = -1;
    char name[256];
    count = 0;    
    if(init())return -1;
    hr = MFCreateAttributes(&pAttributes, 1); if(FAILED(hr))return -1;    
    hr = pAttributes->lpVtbl->SetGUID(pAttributes, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, type); if(FAILED(hr))goto clean;    
    hr = MFEnumDeviceSources(pAttributes, &ppDevices, &count); if(FAILED(hr))goto clean;
    if(count <= 0)goto clean;
    if(!devname){
        hr = ppDevices[0]->lpVtbl->ActivateObject(ppDevices[0], &IID_IMFMediaSource, (LPVOID*)&in->source); if(FAILED(hr))goto clean; 
        if(media_input_init(in))goto clean;
        ret = 0;
    }else{
        for(i = 0; i < count; i++){
            hr = ppDevices[i]->lpVtbl->GetAllocatedString(ppDevices[i], &MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &wname, &len); if(FAILED(hr))goto clean;        
            name[0] = 0;
            WideCharToMultiByte(CP_UTF8, 0, wname, -1, name, 256, NULL, NULL);
            CoTaskMemFree(wname);
            wname = NULL;
            if((devname[0] == '*' && strstr(name, devname + 1)) || strcmp(devname, name) == 0){
                hr = ppDevices[i]->lpVtbl->ActivateObject(ppDevices[i], &IID_IMFMediaSource, (LPVOID*)&in->source); if(FAILED(hr))goto clean;
                if(media_input_init(in))goto clean; 
                ret = 0;
                break;    
            }
        } 
    }

clean:
    for(i = 0; i < count; i++){
        SAFE_RELEASE(ppDevices[i]);
    }    
    SAFE_RELEASE(pAttributes);
    if(ppDevices)CoTaskMemFree(ppDevices);
    if(wname)CoTaskMemFree(wname);    
    return ret;       
}

video_input_t video_input_from_device(const char* name){
    media_input* in = malloc(sizeof(media_input));
    if(!in)return NULL;
    if(media_input_create(in, name, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID)){
        free(in);
        return NULL;
    }
    in->node = NULL;
    return (video_input_t)in;
}

audio_input_t audio_input_from_device(const char* name){
    media_input* in = malloc(sizeof(media_input));
    if(!in)return NULL;
    if(media_input_create(in, name, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_AUDCAP_GUID)){
        free(in);
        return NULL;
    }
    in->node = NULL;
    return (audio_input_t)in;    
}


audio_output_t audio_output_default(){
    HRESULT hr;    
    IMFActivate *pActivate = NULL;
    media_output* out = malloc(sizeof(media_output));
    if(!out)return NULL;
    hr = MFCreateAudioRendererActivate(&pActivate); if(FAILED(hr))goto error;
    out->node = NULL;
    out->act = pActivate;
    return (audio_output_t)out;
error:
    if(out)free(out);
    return NULL;    
}

video_output_t video_output_to_window(void* native_window){
    HRESULT hr;    
    IMFActivate *pActivate = NULL;
    media_output* out = malloc(sizeof(media_output));
    if(!out)return NULL;
    hr = MFCreateVideoRendererActivate((HWND)native_window, &pActivate); if(FAILED(hr))goto error;
    out->node = NULL;
    out->act = pActivate;
    return (video_output_t)out;
error:
    if(out)free(out);
    return NULL;     
}


static int init_control(media_session_t m){
    HRESULT hr;
    IMFVideoDisplayControl* control = NULL;
    if(m->control)return 0;
    hr = MFGetService((IUnknown*)m->session, &MR_VIDEO_RENDER_SERVICE, &IID_IMFVideoDisplayControl, (PVOID*) &control);  if(FAILED(hr))goto error;
    m->control = control;
    return 0;
error:
    return -1;    
}

int media_session_resize_video(media_session_t m, int w, int h){    
    HRESULT hr;
    RECT rcDest = { 0, 0, w, h};
    if(!m)return -1;
    init_control(m);
    if(m->control){
        hr = m->control->lpVtbl->SetVideoPosition(m->control, NULL, &rcDest); if(FAILED(hr))goto error;
    }
    return 0;
error:    
    return -1;    
}

int media_session_set_fullscreen(media_session_t m, bool full){
    HRESULT hr;    
    if(!m)return -1;
    init_control(m);
    if(m->control){
        hr = m->control->lpVtbl->SetFullscreen(m->control, full?TRUE:FALSE); if(FAILED(hr))goto error;
    }
    return 0;
error:    
    return -1;    
}



int video_input_get_size(video_input_t v, int* w, int* h){
    HRESULT hr;
    UINT64 sz;
    int ret = -1;
    media_input* in = (media_input*)v;
    IMFMediaTypeHandler *pHandler = NULL;
    IMFMediaType* pType = NULL;    
    if(!v || !w || !h)return -1;
    hr = in->sd->lpVtbl->GetMediaTypeHandler(in->sd, &pHandler); if (FAILED(hr))goto clean;   
    hr = pHandler->lpVtbl->GetCurrentMediaType(pHandler, &pType); if (FAILED(hr))goto clean;
    hr = pType->lpVtbl->GetUINT64(pType, &MF_MT_FRAME_SIZE, &sz); if (FAILED(hr))goto clean;
    *w = sz >> 32;
    *h = sz & 0xffffffff;
    ret = 0; 
clean:
    SAFE_RELEASE(pType);
    SAFE_RELEASE(pHandler);
    return ret;    
}


static int save_screenshot(BITMAPINFOHEADER* bh, BYTE* pixel, DWORD pixellen, const char* filename){
    HRESULT hr;
    int ret = -1;
    IWICImagingFactory* factory = NULL;
    IWICBitmap* bitmap = NULL;
    IWICBitmapEncoder* encoder = NULL;
    IWICStream* stream = NULL; 
    IWICBitmapFrameEncode* frame = NULL;
    IWICBitmapFlipRotator* rotator = NULL;
    IWICBitmapSource* source;
    WCHAR wname[MAX_PATH];
    int w,h,s;
    WICPixelFormatGUID inputfmt;
    MultiByteToWideChar(CP_UTF8,0,filename,-1,wname,MAX_PATH);
    w = bh->biWidth;
    h = bh->biHeight; 
    if(bh->biCompression != BI_RGB){
        return -2;
    }
    if(bh->biBitCount == 24){
        CopyMemory(&inputfmt, &GUID_WICPixelFormat24bppBGR, sizeof(WICPixelFormatGUID));
        s = w * 3;
        s = s + (4 - (s % 4));
    }else if(bh->biBitCount == 32){        
        CopyMemory(&inputfmt, &GUID_WICPixelFormat32bppBGR, sizeof(WICPixelFormatGUID));
        s = w * 4;
    }else{
        return -3;
    }
    hr = CoCreateInstance(&CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, &IID_IWICImagingFactory, (void**)&factory); if(FAILED(hr))goto clean;    
    if(h > 0){
        hr = factory->lpVtbl->CreateBitmapFlipRotator(factory, &rotator); if(FAILED(hr))goto clean;        
    } else {
        h = 0 - h;
    }    
    hr = factory->lpVtbl->CreateBitmapFromMemory(factory, w, h, &inputfmt, s, pixellen, pixel, &bitmap); if(FAILED(hr))goto clean;
    hr = factory->lpVtbl->CreateEncoder(factory, &GUID_ContainerFormatPng, NULL, &encoder); if(FAILED(hr))goto clean;
    hr = factory->lpVtbl->CreateStream(factory, &stream); if(FAILED(hr))goto clean;  
    hr = stream->lpVtbl->InitializeFromFilename(stream, wname, GENERIC_WRITE); if(FAILED(hr))goto clean;      
    hr = encoder->lpVtbl->Initialize(encoder, (IStream*)stream, WICBitmapEncoderNoCache); if(FAILED(hr))goto clean;
    hr = encoder->lpVtbl->CreateNewFrame(encoder, &frame, NULL); if(FAILED(hr))goto clean;
    hr = bitmap->lpVtbl->GetPixelFormat(bitmap, &inputfmt); if(FAILED(hr))goto clean;
    hr = frame->lpVtbl->Initialize(frame, NULL); if(FAILED(hr))goto clean;      
    hr = frame->lpVtbl->SetSize(frame, w, h); if(FAILED(hr))goto clean;
    hr = frame->lpVtbl->SetPixelFormat(frame, &inputfmt); if(FAILED(hr))goto clean;
    if(rotator){
        hr = rotator->lpVtbl->Initialize(rotator, (IWICBitmapSource*)bitmap, WICBitmapTransformFlipVertical); if(FAILED(hr))goto clean;
        source = (IWICBitmapSource*)rotator;
    }else{
        source = (IWICBitmapSource*)bitmap;          
    }
    hr = frame->lpVtbl->WriteSource(frame, source, NULL); if(FAILED(hr))goto clean; 
    hr = frame->lpVtbl->Commit(frame); if(FAILED(hr))goto clean;
    hr = encoder->lpVtbl->Commit(encoder); if(FAILED(hr))goto clean;    
    ret = 0;
clean:
    SAFE_RELEASE(frame);
    SAFE_RELEASE(encoder);
    SAFE_RELEASE(stream);
    SAFE_RELEASE(rotator);
    SAFE_RELEASE(factory);
    return ret;
}

int media_session_save_screenshot(media_session_t m, const char* filename){
    HRESULT hr;   
    int ret = -1;
    IMFVideoDisplayControl* control;
    MFVideoNormalizedRect rect;
    RECT rc;
    RECT nrc;
    BITMAPINFOHEADER bh;
    BYTE* pixel = NULL;
    DWORD pixellen = 0;
    LONGLONG times = 0; 
    SIZE sz; 
    int szchange = 0;   
    if(!m)return -1;
    init_control(m);
    if(!m->control)return -2;
    control = m->control;
    hr = control->lpVtbl->GetVideoPosition(control, &rect, &rc); if(FAILED(hr))goto clean;
    bh.biSize = sizeof(BITMAPINFOHEADER);
    hr = control->lpVtbl->GetNativeVideoSize(control, &sz, NULL);  if(FAILED(hr))goto clean;
    nrc.left = 0; 
    nrc.right = sz.cx;
    nrc.top = 0;
    nrc.bottom = sz.cy;
    hr = control->lpVtbl->SetVideoPosition(control, &rect, &nrc); if(FAILED(hr))goto clean;
    szchange = 1;
    hr = control->lpVtbl->GetCurrentImage(control, &bh, &pixel, &pixellen, &times); if(FAILED(hr))goto clean;
    ret = save_screenshot(&bh, pixel, pixellen, filename);
clean:
    if(pixel)CoTaskMemFree(pixel);
    if(szchange){
        control->lpVtbl->SetVideoPosition(control, &rect, &rc);
    }
    return ret;        
}



